/******************************************************************************
 * Copyright 2022 The AIR Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

#pragma once
#include <iostream>
#include <memory>
#include <mutex>

#include "base-headers/memory/ring-array.hpp"

namespace base {
namespace memory {

struct SimpleBuffer final {
  std::shared_ptr<void> buffer;
  size_t allocated_length;
  size_t length;
};

class SimpleBufferPool final
    : public std::enable_shared_from_this<SimpleBufferPool> {
 public:
  using ptr_t = std::shared_ptr<SimpleBufferPool>;

  static constexpr size_t kArraySize = 8192;
  static constexpr size_t kBufferSize = 8192;

  static ptr_t NewInstance() {
    return {new SimpleBufferPool(), [](SimpleBufferPool* p) { delete p; }};
  }

  SimpleBuffer Get() {
    std::lock_guard<std::recursive_mutex> lg(lock_);
    void* ptr = nullptr;
    if (0 == array_->ArrayUsed()) {
      if (allocated_buffer_ >= kArraySize) {
        std::cerr << "[Buf Pool] Reach limit!" << std::endl;
        return {};
      }
      ptr = malloc(kBufferSize);
      if (!ptr) {
        std::cerr << "[Buf Pool] Failed to alloc buffer." << std::endl;
        return {};
      }
      allocated_buffer_++;
    } else {
      array_->GetData(&ptr, 1);
    }
    if (!ptr) {
      std::cerr << "[Buf Pool] Internal FATAL ERROR." << std::endl;
      abort();
    }
    bzero(ptr, kBufferSize);
    auto pool = shared_from_this();
    auto res = std::shared_ptr<void>(ptr, [pool](void* p) {
      std::lock_guard<std::recursive_mutex> lg(pool->lock_);
      pool->array_->PutData(&p, 1);
    });
    return {res, kBufferSize, 0};
  }

 private:
  using array_t = RingArray<void*>;
  array_t::ptr_t array_;
  size_t allocated_buffer_;
  std::recursive_mutex lock_{};

  SimpleBufferPool()
      : array_(array_t ::NewInstance(kArraySize,
                                     [](void** p) {
                                       if (p) {
                                         free(*p);
                                       }
                                     })),
        allocated_buffer_(0) {}
  ~SimpleBufferPool() = default;
};

}  // namespace memory
}  // namespace base
